# -*- encoding: utf-8 -*-
'''
@File    :   paderborn_reader.py
@Time    :   2022/3/23 9:01
@Author  :   ZhangChaoYang
@Desc    :   把Paderborn原始数据文件，经过各种转换（FFT、CWT等）变成numpy ndarray，存入文件。德国帕德伯恩大学轴承数据集 http://groups.uni-paderborn.de/kat/BearingDataCenter/。由加速寿命测试引起的真实损坏：KA04,KA15,KB23,KB24,KI14,KI16
'''

import os
import sys

sys.path.insert(0, os.getcwd())
import numpy as np

from data_process import mat_reader
from data_process.mat_reader import convertDir1d, convertDir2d, take_a_look_1d, take_a_look_2d
from absl import flags
from absl import app
import random

FLAGS = flags.FLAGS

flags.DEFINE_string('source_data', "E:\zcy\数据集\Paderborn", '原始数据存在目录')

fs = 64e3  # 采样频率64kHz
WC = ["N15_M07_F10", "N09_M07_F10", "N15_M01_F10",
      "N15_M07_F04"]  # work condition


class PUReader(mat_reader.MatFileReader):
    def __init__(self, infile, trans1d="", trans2d="", work_load=""):
        super(PUReader, self).__init__(
            infile, "paderborn", trans1d, trans2d, work_load)

    def window_size(self):
        return 4096  # 在paderborn数据集上，4266个采样点才能保证拿到一转的数据

    def get_label(self):
        paths = os.path.split(os.path.dirname(self.infile))
        if paths[-1].startswith("K0"):  # 最后一级目录以"K0"开头为正常样本
            return 1
        else:
            return 0

    def read1d(self):
        if random.random() < 0.5:  # TODO PU数据集太大，随机丢弃一半数据
            return None
        dic = self.read_file_to_dict()
        basename = self.get_basename()
        if self.work_load:
            arr = basename.split("_", 4)
            if len(arr) <= 4 or "_".join(arr[:3]) != self.work_load:
                return None
        # print(dic.keys())#['__header__', '__version__', '__globals__', 'N09_M07_F10_K001_1']
        for key in dic.keys():
            if key == basename[:-4]:
                array = np.transpose(
                    dic[key][0][0][2][0][6][2][0])  # 振动数据vibration_1
                x = self.trans1d(array, 1, fs)
                return x

    def read2d(self):
        if random.random() < 0.5:  # TODO PU数据集太大，随机丢弃一半数据
            return None
        dic = self.read_file_to_dict()
        basename = self.get_basename()
        if self.work_load:
            arr = basename.split("_", 4)
            if len(arr) <= 4 or "_".join(arr[:3]) != self.work_load:
                return None
        for key in dic.keys():
            if key == basename[:-4]:
                array = np.transpose(
                    dic[key][0][0][2][0][6][2][0])  # 振动数据vibration_1
                x = self.trans2d(array, 1, fs)
                return x


def main(argv):
    # for trans1d in ['original', 'stat', 'fft']:
    for trans1d in ['original']:
        reader = PUReader("", trans1d=trans1d)
        convertDir1d(reader, [FLAGS.source_data])
        for wc in WC:
            reader = PUReader("", trans1d=trans1d, work_load=wc)
            convertDir1d(reader, [FLAGS.source_data])
        take_a_look_1d(reader.name, trans1d)

    for trans2d in ['sfft', 'cwt']:
        reader = PUReader(FLAGS.source_data, trans2d=trans2d)
        convertDir2d(reader, [FLAGS.source_data])
        for wc in WC:
            reader = PUReader("", trans2d=trans2d, work_load=wc)
            convertDir2d(reader, [FLAGS.source_data])
        take_a_look_2d(reader.name, trans2d)


if __name__ == '__main__':
    app.run(main)

# python .\data_process\paderborn_reader.py --source_data E:\zcy\数据集\Paderborn
